import numpy as np
import paddle
from paddle.nn import functional as F


class Rect(object):
    """
        点的生成方式支持两种类型：
        - bound，`position`表示为(xmin, ymin, xmax, ymax)
        - center, `position`表示为(xcenter, ycenter, width, height)
        """

    def __init__(self, position, mode='bound'):
        if mode == 'bound':
            self.xmin, self.ymin, self.xmax, self.ymax = position
        else:
            self.xmin = position[0] - position[2] / 2
            self.ymin = position[1] - position[3] / 2
            self.xmax = self.xmin + position[2]
            self.ymax = self.ymin + position[3]

    @property
    def area(self):
        return (self.xmax - self.xmin) * (self.ymax - self.ymin)


def batch_iou(rect: np.ndarray, to_rects_batch: np.ndarray) -> np.ndarray:
    """
        计算rect与to_rects_batch中所有矩形的iou

        Args:
        ---
            rect: 单一矩形，[[xmin, ymin, xmax, ymax]]
            to_rects_batch: 批量矩形，[[xmin, ymin, xmax, ymax], ...]
    """
    xmin = np.maximum(to_rects_batch[:, 0], rect[:, 0])
    ymin = np.maximum(to_rects_batch[:, 1], rect[:, 1])
    xmax = np.minimum(to_rects_batch[:, 2], rect[:, 2])
    ymax = np.minimum(to_rects_batch[:, 3], rect[:, 3])
    inter_h = np.clip(xmax - xmin, 0, None)
    inter_w = np.clip(ymax - ymin, 0, None)
    inter_area = inter_w * inter_h

    a1 = (rect[:, 2] - rect[:, 0]) * (rect[:, 3] - rect[:, 1])
    a2 = (to_rects_batch[:, 2] - to_rects_batch[:, 0]) \
        * (to_rects_batch[:, 3] - to_rects_batch[:, 1])
    return inter_area / (a1 + a2 - inter_area)


def center2boundary(xcenter, ycenter, height, width):
    return xcenter - height / 2, ycenter - width / 2, xcenter + height / 2, ycenter + width / 2


def get_objectness_label(img, label_rects, anchors, downsamples=32) -> np.ndarray:
    """
    标注数据，产生(tx, ty, th, tw, objectness)的预测框。

    Args:
    ---
        :param img: 输入图片。

    Return:
    ---
        label: [`batch_size`, `num_anchors`, `num_rows`, `num_cols`, 5]
    """
    batch_size, h, w, channels = img.shape
    num_anchors = len(anchors)
    num_rows = h // downsamples
    num_cols = w // downsamples
    label = np.zeros((batch_size, num_anchors, num_rows, num_cols, 5))
    for batch_id in range(batch_size):
        for gt_label in label_rects[batch_id]:
            # gt_label [xmin, ymin, xmax, ymax]
            xmin, ymin, xmax, ymax = gt_label[0], gt_label[1], gt_label[2], gt_label[3]
            # xcenter after downsamples
            xcenter = (xmax + xmin) / 2 / downsamples
            # ycenter after downsamples
            ycenter = (ymax + ymin) / 2 / downsamples
            # height after downsamples
            height = (xmax - xmin) / downsamples
            # width after down samples
            width = (ymax - ymin) / downsamples

            if min(height, width) < 1e-10:
                continue
            # left top point of center
            i = int(xcenter)
            j = int(ycenter)

            anchor_rects = []

            for anchor in anchors:
                batch_rect = center2boundary(
                    xcenter, ycenter, anchor[1] / downsamples, anchor[0] / downsamples)
                anchor_rects.append(list(batch_rect))
            anchor_rects = np.array(anchor_rects, dtype='float32')
            iou_results = batch_iou(
                np.array([[xmin, ymin, xmax, ymax]]), anchor_rects)
            sorted_ids = iou_results.argsort()
            ka = sorted_ids[-1]
            # objectness
            label[batch_id, ka, i, j, 4] = 1
            # x bias
            label[batch_id, ka, i, j, 0] = xcenter - i
            # y bias
            label[batch_id, ka, i, j, 1] = ycenter - j
            # height scale
            label[batch_id, ka, i, j, 2] = np.log(
                height / (anchors[ka][1] / downsamples))
            # width scale
            label[batch_id, ka, i, j, 3] = np.log(
                width / (anchors[ka][0] / downsamples))
            # TODO scale location
    # assert isinstance(label, np.ndarray)
    return label.transpose((0, 1, 4, 2, 3))

def xy2yx(point):
    """
    将坐标轴表示互换
    """
    return [point[1], point[0]]

def predict(output, anchors, downsamples=32):
    
    batch_size, channels, num_rows, num_cols = output.shape
    num_anchors = int(channels / 5)
    prediction = []
    # output: [batch_size, 3 * 5, 20, 10]
    output = paddle.transpose(output, (0, 2, 3, 1))
    output = paddle.reshape(output, (batch_size, num_rows, num_cols, num_anchors, 5))
    output[:, :, :, :, 4] = F.sigmoid(output[:, :, :, :, 4])
    output[:, :, :, :, :2] = F.sigmoid(output[:, :, :, :, :2])
    output = output.numpy()
    for batch_id in range(batch_size):
        pred_rects = []
        for i in range(num_rows):
            for j in range(num_cols):
                for k in range(num_anchors):
                    if output[batch_id, i, j, k, 4] > 0.5:
                        xcenter = i + output[batch_id, i, j, k, 0]
                        ycenter = j + output[batch_id, i, j, k, 1]
                        height = np.exp(anchors[k][1] / downsamples * output[batch_id, i, j, k, 2])
                        width = np.exp(anchors[k][0] / downsamples * output[batch_id, i, j, k, 3])
                        print(height, width)
                        # 生成四个点的坐标
                        xmin = np.clip(int((xcenter - height / 2) * downsamples), 0, (num_rows + 1) * downsamples - 1)
                        ymin = np.clip(int((ycenter - width / 2) * downsamples), 0, (num_cols + 1) * downsamples - 1)
                        xmax = np.clip(int((xcenter + height / 2) * downsamples), 0, (num_rows + 1) * downsamples - 1)
                        ymax = np.clip(int((ycenter + width / 2) * downsamples), 0, (num_cols + 1) * downsamples - 1)
                        pred_rects.append([xmin, ymin, xmax, ymax])
        prediction.append(pred_rects)
    return prediction

def prediction_visualize(img, output, anchors, downsamples=32):
    """
    img, output都是批量的
    img要求是 NHWC
    """
    prediction = predict(output, anchors, downsamples)
    batch_size = output.shape[0]
    for batch_id in range(batch_size):
        new_img = img[batch_id]
        for rect in prediction[batch_id]:
            new_img = cv2.rectangle(new_img, (rect[1], rect[0]), (rect[3], rect[2]), (0, 255, 255), 2)
        cv2.imshow(str(batch_id), new_img)
        cv2.waitKey(0)
        cv2.destroyAllWindows()
    return prediction

if __name__ == '__main__':
    import cv2
    import pandas

    from transform import random_transform

    label = [197,250,283,510]
    img = cv2.imread('./data/training_images/vid_4_700.jpg')
    new_img = cv2.rectangle(img, (label[1], label[0]), (label[3], label[2]), color=(0, 0, 255), thickness=2)
    cv2.imshow('win', new_img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
    img, label = random_transform(img, labels=[label])
    label = label[0]
    new_img = cv2.rectangle(img, (label[1], label[0]), (label[3], label[2]), color=(0, 255, 255), thickness=2)
    anchors = [
        [16, 8],
        [64, 32],
        [128, 64]
    ]
    predict_label = get_objectness_label(np.array([img]), np.array([[label]]), anchors=anchors)
    predict_label = predict_label[0]
    #             print(ka, anchor_rects[2])
    # 2 [ 6.1875   9.03125  8.1875  13.03125]
    new_img = cv2.rectangle(new_img, [int(9.03125 * 32)-64, int(6.1875 * 32)-32], [int(32 * 13.03125)+64, int(32 * 8.1875)+32], color=(255, 0, 0), thickness=1)
    cv2.imshow('win', new_img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

